package com.archgeek.java.udp;

import java.io.IOException;
import java.net.DatagramPacket;
import java.net.DatagramSocket;
import java.net.InetAddress;
import java.net.InetSocketAddress;
import java.net.UnknownHostException;

/**
 * @author pizhihui
 * @date 2023-08-07 21:32
 */
public class TaskManagerDiscoverService {


    public static void main(String[] args) throws IOException {


        int TM_PORT = 6123;
        int magicNumber = 0;
        int discoveryPort = 7001;

        InetAddress jobManagerAddress = new InetSocketAddress("127.0.0.1", TM_PORT).getAddress();

        DatagramSocket socket = new DatagramSocket();

        final DatagramPacket addressRequest = createTaskManagerAddressRequestPacket(magicNumber);
        addressRequest.setAddress(jobManagerAddress);
        addressRequest.setPort(TM_PORT);

        socket.send(addressRequest);

        final DatagramPacket responsePacket = new DatagramPacket(new byte[RESPONSE_PACKET_SIZE], RESPONSE_PACKET_SIZE);
        socket.receive(responsePacket);
        InetAddress taskManagerAddress = extractInetAddress(responsePacket);
        System.out.println(responsePacket.getSocketAddress());
        System.out.println(responsePacket.getAddress());
        System.out.println(responsePacket.getPort());
        System.out.println("tm address: " + taskManagerAddress + " port: " + extractIpcPort(responsePacket));

        System.in.read();


    }

    private static final int RESPONSE_PACKET_SIZE = 64;

    private static final int TM_ADDRESS_REQUEST_ID = 2;


    private static final int MAGIC_NUMBER_OFFSET = 0;
    private static final int PACKET_ID_OFFSET = 4;
    private static final int PACKET_TYPE_ID_OFFSET = 8;

    private static final int PAYLOAD_OFFSET = 12;


    private static DatagramPacket createTaskManagerAddressRequestPacket(final int magicNumber) {

        final byte[] bytes = new byte[12];
        integerToByteArray(magicNumber, MAGIC_NUMBER_OFFSET, bytes);
        integerToByteArray(generateRandomPacketID(), PACKET_ID_OFFSET, bytes);
        integerToByteArray(TM_ADDRESS_REQUEST_ID, PACKET_TYPE_ID_OFFSET, bytes);

        return new DatagramPacket(bytes, bytes.length);
    }

    private static int generateRandomPacketID() {

        return (int) (Math.random() * (double) Integer.MAX_VALUE);
    }

    private static void integerToByteArray(final int integerToSerialize, final int offset, final byte[] byteArray) {

        for (int i = 0; i < 4; ++i) {
            final int shift = i << 3; // i * 8
            byteArray[(offset + 3) - i] = (byte) ((integerToSerialize & (0xff << shift)) >>> shift);
        }
    }

    private static InetAddress extractInetAddress(DatagramPacket packet) {

        final byte[] data = packet.getData();

        if (data == null) {
            return null;
        }

        if (packet.getLength() < PAYLOAD_OFFSET + 8) {
            return null;
        }

        final int len = byteArrayToInteger(data, PAYLOAD_OFFSET);

        final byte[] addr = new byte[len];
        System.arraycopy(data, PAYLOAD_OFFSET + 4, addr, 0, len);

        InetAddress inetAddress = null;

        try {
            inetAddress = InetAddress.getByAddress(addr);
        } catch (UnknownHostException e) {
            return null;
        }

        return inetAddress;
    }

    private static int extractIpcPort(DatagramPacket packet) {

        final byte[] data = packet.getData();

        if (data == null) {
            return -1;
        }

        if (packet.getLength() < (PAYLOAD_OFFSET + 4)) {
            return -1;
        }

        return byteArrayToInteger(data, PAYLOAD_OFFSET);
    }

    private static int byteArrayToInteger(final byte[] byteArray, final int offset) {

        int integer = 0;

        for (int i = 0; i < 4; ++i) {
            integer |= (byteArray[(offset + 3) - i] & 0xff) << (i << 3);
        }

        return integer;
    }

}
